import numpy as np
from typing import Union
import platgo as pg


class DE:
    """
    差分变异算子类

    当传入的是一个种群时，对整个种群做变异
    当传入的是3个一维矩阵时，视为3个个体，对个体做变异
    """

    # TODO: 需要增加其它差分算子

    def __init__(self, F: float = 0.5) -> None:
        self.F = F

    def _mut_individual(self, dec1: np.ndarray, dec2: np.ndarray, dec3: np.ndarray) -> np.ndarray:
        """
        对单个个体变异，参数是数组类型
        :param dec1:
        :param dec2:
        :param dec3:
        :return:
        """
        return dec1 + self.F * (dec2 - dec3)

    def _mut_pop(self, pop: pg.Population) -> pg.Population:
        """
        对种群内所有个体变异
        :param pop:
        :return:
        """
        decs1 = pop.decs.copy()
        decs2 = pop.decs.copy()
        decs3 = pop.decs.copy()
        np.random.shuffle(decs2)
        np.random.shuffle(decs3)
        new_decs = decs1 + self.F * (decs2 - decs3)
        new_pop = pg.Population(decs=new_decs)
        return new_pop

    def _mut_three_pop(self, pop1: pg.Population, pop2: pg.Population, pop3: pg.Population) -> pg.Population:
        decs1 = pop1.decs
        decs2 = pop2.decs
        decs3 = pop3.decs
        new_decs = decs1 + self.F * (decs2 - decs3)
        new_pop = pg.Population(decs=new_decs)
        return new_pop

    def __call__(self, p1: Union[pg.Population, np.ndarray], p2: Union[pg.Population, np.ndarray] = None, p3: Union
    [pg.Population, np.ndarray] = None) -> Union[pg.Population, np.ndarray]:
        if p2 is None:  # mutate for a population
            return self._mut_pop(p1)
        elif isinstance(p2, np.ndarray):  # mutate for a individual
            return self._mut_individual(p1, p2, p3)
        else:  # mutate for three specified population
            return self._mut_three_pop(p1, p2, p3)
